# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold, cross_validate
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, roc_auc_score, accuracy_score, f1_score, roc_curve
import logging
import os
import matplotlib.pyplot as plt  # 新增
import seaborn as sns


# ==================== 日志类 ====================
class Logger(object):
    level_relations = {
        'debug': logging.DEBUG,
        'info': logging.INFO,
        'warning': logging.WARNING,
        'error': logging.ERROR
    }

    def __init__(self, root_path, log_name, level='info', fmt='%(asctime)s - %(levelname)s: %(message)s'):
        self.root_path = root_path
        self.log_name = log_name
        self.fmt = fmt
        self.logger = logging.getLogger(log_name)
        if not self.logger.handlers:
            self.logger.setLevel(self.level_relations.get(level))

    def get_logger(self):
        path = os.path.join(self.root_path, 'log')
        os.makedirs(path, exist_ok=True)
        file_name = os.path.join(path, self.log_name + '.log')
        fileHandler = logging.FileHandler(file_name, encoding="utf-8", mode="a")
        formatter = logging.Formatter(self.fmt)
        fileHandler.setFormatter(formatter)
        if not any(isinstance(h, logging.FileHandler) for h in self.logger.handlers):
            self.logger.addHandler(fileHandler)
        return self.logger


# ==================== 数据与模型部分 ====================
TRAIN_PATH = "../data/train.csv"
TEST_PATH = "../data/test.csv"
TARGET = "Attrition"


def load_data():
    train = pd.read_csv(TRAIN_PATH)
    test = pd.read_csv(TEST_PATH)

    drop_cols = []
    for col in ["EmployeeNumber", "Over18", "StandardHours"]:
        if col in train.columns and (col == "EmployeeNumber" or train[col].nunique(dropna=False) <= 1):
            drop_cols.append(col)

    X = train.drop(columns=[TARGET] + drop_cols, errors="ignore")
    y = train[TARGET].astype(int)

    X_test = test.drop(columns=[TARGET] + drop_cols, errors="ignore")
    y_test = test[TARGET].astype(int) if TARGET in test.columns else None

    return X, y, X_test, y_test, test, drop_cols


# 添加一个函数用于绘制特征分析图
def plot_feature_analysis(X, y, result_dir="../data/fig"):
    os.makedirs(result_dir, exist_ok=True)
    X = X.copy()
    X['Target'] = y.values

    # 选择一些关键特征进行绘制（包括原始特征和组合特征）
    num_features = ['Age', 'MonthlyIncome', 'JobSatisfaction', 'EnvironmentSatisfaction',
                    'Age_YearsAtCompany', 'Stability', 'YearsAtCompany', 'TotalWorkingYears']
    cat_features = ['Department', 'JobRole', 'OverTime', 'Age_bin', 'Income_bin']

    # 只绘制存在于X中的特征
    num_features = [f for f in num_features if f in X.columns]
    cat_features = [f for f in cat_features if f in X.columns]

    # 绘制数值特征与目标变量的箱线图
    for feature in num_features:
        if feature in X.columns:
            plt.figure(figsize=(10, 6))
            sns.boxplot(x='Target', y=feature, data=X)
            plt.title(f'Boxplot of {feature} by Attrition')
            save_path = os.path.join(result_dir, f"boxplot_{feature}.png")
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            print(f"箱线图已保存至: {save_path}")

    # 绘制分类特征与目标变量的条形图（显示Attrition率）
    for feature in cat_features:
        if feature in X.columns:
            plt.figure(figsize=(10, 6))
            # 计算每个类别的Attrition率
            attrition_rate = X.groupby(feature)['Target'].mean().sort_values(ascending=False)
            attrition_rate.plot(kind='bar')
            plt.title(f'Attrition Rate by {feature}')
            plt.ylabel('Attrition Rate')
            save_path = os.path.join(result_dir, f"barplot_{feature}.png")
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            print(f"条形图已保存至: {save_path}")

    # 绘制数值特征的相关性热力图
    plt.figure(figsize=(12, 10))
    corr = X[num_features].corr()
    sns.heatmap(corr, annot=True, cmap='coolwarm', center=0)
    plt.title('Correlation Heatmap of Numerical Features')
    save_path = os.path.join(result_dir, "correlation_heatmap.png")
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"相关性热力图已保存至: {save_path}")

def add_interactions(X):
    X = X.copy()
    if 'Age' in X.columns and 'YearsAtCompany' in X.columns:
        X['Age_YearsAtCompany'] = np.exp(X['Age']) * np.exp(X['YearsAtCompany'])
    if 'Department' in X.columns and 'JobRole' in X.columns:
        X['Dept_JobRole'] = X['Department'].astype(str) + '_' + X['JobRole'].astype(str)
    if 'JobLevel' in X.columns and 'MonthlyIncome' in X.columns:
        X['JobLevel_MonthlyIncome'] = np.exp(X['JobLevel']) + X['MonthlyIncome']
    if 'YearsAtCompany' in X.columns and 'TotalWorkingYears' in X.columns:
        X['YearsAtCompany_per_Total'] = X['YearsAtCompany'] / (X['TotalWorkingYears'] + 5)
    if 'Age' in X.columns and 'YearsSinceLastPromotion' in X.columns:
        X['Age_YearsSinceLastPromotion'] = X['Age'] * X['YearsSinceLastPromotion']
    if 'YearsInCurrentRole' in X.columns and 'YearsWithCurrManager' in X.columns:
        X['YearsInCurrentRole_YearsWithCurrManager'] = X['YearsInCurrentRole'] * X['YearsWithCurrManager']
    if 'MonthlyIncome' in X.columns and 'JobSatisfaction' in X.columns:
        X['Income_JobSatisfaction'] = X['MonthlyIncome'] * X['JobSatisfaction']
    if 'MonthlyIncome' in X.columns and 'EnvironmentSatisfaction' in X.columns:
        X['Income_EnvironmentSatisfaction'] = X['MonthlyIncome'] * X['EnvironmentSatisfaction']
    if 'YearsAtCompany' in X.columns and 'OverTime' in X.columns:
        X['Overtime_YearsAtCompany'] = X['YearsAtCompany'] * X['OverTime'].map({'Yes': 1, 'No': 0})
    if 'YearsAtCompany' in X.columns and 'NumCompaniesWorked' in X.columns:
        X['Stability'] = X['YearsAtCompany'] / (X['NumCompaniesWorked'] + 1)
    if 'Age' in X.columns:
        X['Age_bin'] = pd.cut(X['Age'], bins=[17, 25, 35, 45, 60], labels=['<25', '25-35', '35-45', '45+'])
        X['Age_bin'] = X['Age_bin'].astype(str)
    if 'MonthlyIncome' in X.columns:
        X['Income_bin'] = pd.qcut(X['MonthlyIncome'], q=4, labels=['low', 'med_low', 'med_high', 'high'])
        X['Income_bin'] = X['Income_bin'].astype(str)
    return X


def build_preprocessor(X):
    cat_cols = [c for c in X.columns if (X[c].dtype == "object" or str(X[c].dtype) == "category")]
    num_cols = [c for c in X.columns if c not in cat_cols]

    numeric_transformer = Pipeline([
        ("imputer", SimpleImputer(strategy="median")),
        ("scaler", StandardScaler())
    ])

    categorical_transformer = Pipeline([
        ("imputer", SimpleImputer(strategy="most_frequent")),
        ("onehot", OneHotEncoder(handle_unknown="ignore"))
    ])

    preprocess = ColumnTransformer([
        ("num", numeric_transformer, num_cols),
        ("cat", categorical_transformer, cat_cols)
    ], remainder="drop")

    return preprocess


def build_logistic_regression_pipeline(X):
    preprocess = build_preprocessor(X)
    clf = LogisticRegression(max_iter=2000, class_weight={1: 2.4}, C=0.9)
    pipe = Pipeline([
        ("preprocess", preprocess),
        ("model", clf)
    ])
    return pipe


def cv_evaluate(pipe, X, y):
    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    scores = cross_validate(
        pipe, X, y,
        scoring={"roc_auc": "roc_auc", "accuracy": "accuracy", "f1": "f1"},
        cv=cv, n_jobs=1, return_train_score=False
    )

    summary = {
        "roc_auc_mean": scores["test_roc_auc"].mean(),
        "roc_auc_std": scores["test_roc_auc"].std(),
        "acc_mean": scores["test_accuracy"].mean(),
        "acc_std": scores["test_accuracy"].std(),
        "f1_mean": scores["test_f1"].mean(),
        "f1_std": scores["test_f1"].std(),
    }
    return summary


def evaluate(pipe, X_test, y_test):
    proba = pipe.predict_proba(X_test)[:, 1]
    pred = (proba >= 0.8).astype(int)

    # 打印指标
    print("\n[TEST] ROC-AUC: %.4f" % roc_auc_score(y_test, proba))
    print("[TEST] ACC: %.4f" % accuracy_score(y_test, pred))
    print("[TEST] F1: %.4f" % f1_score(y_test, pred))
    print("\n[TEST] Classification Report:")
    print(classification_report(y_test, pred, digits=4))

    # === 新增：绘制并保存 ROC 曲线 ===
    fpr, tpr, _ = roc_curve(y_test, proba)
    auc_score = roc_auc_score(y_test, proba)

    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {auc_score:.4f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=1, linestyle='--', label='Random Classifier')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic (ROC) Curve')
    plt.legend(loc="lower right")
    plt.grid(alpha=0.3)

    # 保存图像
    result_dir = "../data/fig"
    os.makedirs(result_dir, exist_ok=True)
    save_path = os.path.join(result_dir, "roc_curve.png")
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

    print(f"\nROC 曲线已保存至: {save_path}")


# ==================== 主函数（带日志） ====================
def main():
    logger = Logger(root_path="../log", log_name="model_training", level='info').get_logger()
    logger.info("=" * 50)
    logger.info("开始执行逻辑回归模型训练与评估流程")

    try:
        logger.info("步骤 1: 加载训练与测试数据...")
        X, y, X_test, y_test, *_ = load_data()
        logger.info(f"训练集形状: {X.shape}, 测试集形状: {X_test.shape}")

        logger.info("步骤 2: 添加交互特征与分箱特征...")
        X = add_interactions(X)
        X_test = add_interactions(X_test)
        logger.info(f"特征工程后，训练集特征数: {X.shape[1]}")

        # 在模型训练前，先绘制特征分析图
        logger.info("步骤 3: 绘制特征分析图...")
        plot_feature_analysis(X, y)

        logger.info("步骤 4: 构建逻辑回归管道...")
        pipe = build_logistic_regression_pipeline(X)
        logger.info("管道构建成功")

        logger.info("步骤5: 在全量训练集上拟合模型并评估测试集...")
        pipe.fit(X, y)
        evaluate(pipe, X_test, y_test)
        logger.info("模型训练与评估流程顺利完成！")

    except Exception as e:
        logger.error("程序运行过程中发生异常！", exc_info=True)
        raise


if __name__ == "__main__":
    main()